import numpy as np

from traits.api import Range, Int, Bool, Array, Instance, Enum, on_trait_change, Button
from traitsui.api import View, Item, Tabbed, HGroup, VGroup, VSplit, EnumEditor, TextEditor

import logging

import hardware.api as ha

from tools.emod import ManagedJob

from tools.utility import GetSetItemsMixin

"""
Several options to decide when to start  and when to restart a job, i.e. when to clear data, etc.

1. set a 'new' flag on every submit button

pro: simple, need not to think about anything in subclass

con: continue of measurement only possible by hack (manual submit to JobManager without submit button)
     submit button does not do what it says
     
2. check at start time whether this is a new measurement.

pro: 

con: complicated checking needed
     checking has to be reimplemented on sub classes
     no explicit way to restart the same measurement

3. provide user settable clear / keep flag

pro: explicit

con: user can forget

4. provide two different submit buttons: submit, resubmit

pro: explicit

con: two buttons that user may not understand
     user may use wrong button
     wrong button can result in errors

"""

# utility function
def find_laser_pulses(sequence):
    n = 0
    for channels, t in sequence:
        if 'laser' in channels:
            n+=1
    return n


class Pulsed( ManagedJob, GetSetItemsMixin ):
    
    """Defines a pulsed measurement."""
    
    keep_data = Bool(False) # helper variable to decide whether to keep existing data

    resubmit_button = Button(label='resubmit', desc='Submits the measurement to the job manager. Tries to keep previously acquired data. Behaves like a normal submit if sequence or time bins have changed since previous run.')

    sequence = Instance( list, factory=list )
    
    record_length = Range(low=100, high=100000., value=3000, desc='length of acquisition record [ns]', label='record length [ns]', mode='text', auto_set=False, enter_set=True)
    bin_width = Range(low=0.1, high=1000., value=1.0, desc='bin width [ns]', label='bin width [ns]', mode='text', auto_set=False, enter_set=True)
    
    n_laser = Int(2)
    n_bins = Int(2)
    time_bins = Array(value=np.array((0,1)))
    
    count_data = Array( value=np.zeros((2,2)) )
    
    def submit(self):
        """Submit the job to the JobManager."""
        self.keep_data = False
        ManagedJob.submit(self)

    def resubmit(self):
        """Submit the job to the JobManager."""
        self.keep_data = True
        ManagedJob.submit(self)

    def _resubmit_button_fired(self):
        """React to start button. Submit the Job."""
        self.resubmit() 

    def generate_sequence(self):
        return []

    def apply_parameters(self):
        """Apply the current parameters and decide whether to keep previous data."""
        n_bins = int(self.record_length / self.bin_width)
        time_bins = self.bin_width*np.arange(n_bins)
        sequence = self.generate_sequence()
        n_laser = find_laser_pulses(sequence)

        if self.keep_data and sequence == self.sequence and np.all(time_bins == self.time_bins): # if the sequence and time_bins are the same as previous, keep existing data
            self.old_count_data = self.count_data.copy()
        else:
            self.old_count_data = np.zeros((n_laser,n_bins))
        
        self.sequence = sequence 
        self.time_bins = time_bins
        self.n_bins = n_bins
        self.n_laser = n_laser
        self.keep_data = True # when job manager stops and starts the job, data should be kept. Only new submission should clear data.

    def start_up(self):
        """Put here additional stuff to be executed at startup."""
        pass

    def shut_down(self):
        """Put here additional stuff to be executed at shut_down."""
        pass

    def _run(self):
        """Acquire data."""

        try: # try to run the acquisition from start_up to shut_down
            self.state='run'
            self.apply_parameters()
            self.start_up()
            ha.PulseGenerator().Night()
            tagger = ha.TimeTagger.Pulsed(self.n_bins, int(np.round(self.bin_width*1000)), self.n_laser, 0, 2, 3)
            ha.PulseGenerator().Sequence(self.sequence)

            while True:
                self.thread.stop_request.wait(1.0)
                if self.thread.stop_request.isSet():
                    break
                self.count_data = self.old_count_data + tagger.getData()

            del tagger
            self.shut_down()
            ha.PulseGenerator().Light()

        finally: # if anything fails, recover
            self.state='idle'

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              HGroup(Item('bin_width', width=-80, enabled_when='state == "idle"'),
                                     Item('record_length', width=-80, enabled_when='state == "idle"'),
                                     ),
                              ),
                       title='Pulsed Measurement',
                       )

    get_set_items = ['__doc__','record_length','bin_width','n_bins','time_bins','n_laser','sequence','count_data']
    get_set_order = ['time_bins','count_data']


class Rabi( Pulsed ):
    
    """Defines a Rabi measurement."""

    frequency   = Range(low=1,      high=20e9,  value=2.8705e9, desc='microwave frequency', label='frequency [Hz]', mode='text', auto_set=False, enter_set=True)
    power       = Range(low=-100.,  high=25.,   value=-20,      desc='microwave power',     label='power [dBm]',    mode='text', auto_set=False, enter_set=True)
    switch      = Enum( 'A', 'B',   desc='switch to use for microwave pulses',     label='switch' )


    tau_begin   = Range(low=0., high=1e8,       value=0.,     desc='tau begin [ns]',  label='tau begin [ns]',   mode='text', auto_set=False, enter_set=True)
    tau_end     = Range(low=1., high=1e8,       value=300.,     desc='tau end [ns]',    label='tau end [ns]',     mode='text', auto_set=False, enter_set=True)
    tau_delta   = Range(low=1., high=1e6,       value=3.,      desc='delta tau [ns]',  label='delta tau [ns]',   mode='text', auto_set=False, enter_set=True)
    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)
    wait        = Range(low=1., high=100000.,   value=1000.,    desc='wait [ns]',       label='wait [ns]',        mode='text', auto_set=False, enter_set=True)

    tau = Array( value=np.array((0.,1.)) )

    def apply_parameters(self):
        """Overwrites apply_parameters() from pulsed. Prior to generating sequence, etc., generate the tau mesh."""
        self.tau = np.arange(self.tau_begin, self.tau_end, self.tau_delta)
        Pulsed.apply_parameters(self)
        
    def start_up(self):
        ha.PulseGenerator().Night()
        ha.Microwave().setOutput(self.power, self.frequency)

    def shut_down(self):
        ha.PulseGenerator().Light()
        ha.Microwave().setOutput(None, self.frequency)

    def generate_sequence(self):
        if self.switch == 'A':
            MW = 'mw_a'
        else:
            MW = 'mw_b'
        tau = self.tau
        laser = self.laser
        wait = self.wait
        sequence = []
        for t in tau:
            sequence += [  ([MW],t),  (['laser','aom'],laser),  ([],wait)  ]
        sequence += [ (['sequence'], 100  )  ]
        return sequence

    get_set_items = Pulsed.get_set_items + ['frequency','power','switch','tau_begin','tau_end','tau_delta','laser','wait','tau']
    get_set_order = ['tau','time_bins','count_data']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),
                                                   Item('switch', style='custom', editor=EnumEditor(cols=3, values={'A':'1:A', 'B':'2:B'}),),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                                   label='settings'),
                              ),
                        ),
                       title='Rabi Measurement',
                  )


class T1( Pulsed ):
    
    """Defines a T1 measurement."""

    tau_begin   = Range(low=1., high=1e8,       value=10.0,     desc='tau begin [ns]',  label='tau begin [ns]',   mode='text', auto_set=False, enter_set=True)
    tau_end     = Range(low=1., high=1e8,       value=500.,     desc='tau end [ns]',    label='tau end [ns]',     mode='text', auto_set=False, enter_set=True)
    tau_delta   = Range(low=1., high=1e6,       value=50.,      desc='delta tau [ns]',  label='delta tau [ns]',   mode='text', auto_set=False, enter_set=True)
    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)

    tau = Array( value=np.array((0.,1.)) )

    def apply_parameters(self):
        """Overwrites apply_parameters() from pulsed. Prior to generating sequence, etc., generate the tau mesh."""
        self.tau = np.arange(self.tau_begin, self.tau_end, self.tau_delta)
        Pulsed.apply_parameters(self)
        
    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        sequence = []
        for t in tau:
            sequence += [  ([],t),  (['laser','aom'],laser)  ]
        sequence += [  ([],1000)  ]
        sequence += [  (['sequence'],100)  ]
        return sequence

    get_set_items = Pulsed.get_set_items + ['tau_begin','tau_end','tau_delta','laser','tau']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              ),
                              Tabbed(VGroup(HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                                   label='settings'),
                              ),
                       title='T1 Measurement',
                       )



class PulseCal( Pulsed ):
    
    """Pulse Calibration."""

    frequency   = Range(low=1,      high=20e9,  value=2.8705e9, desc='microwave frequency', label='frequency [Hz]', mode='text', auto_set=False, enter_set=True)
    power       = Range(low=-100.,  high=25.,   value=-20,      desc='microwave power',     label='power [dBm]',    mode='text', auto_set=False, enter_set=True)
    switch      = Enum( 'mw_x', 'mw_y',   desc='switch to use for microwave pulses',     label='switch' )
    #source      = Enum( 'A', 'B',   desc='microwave source', label='source' )

    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)
    wait        = Range(low=1., high=100000.,   value=1000.,    desc='wait [ns]',       label='wait [ns]',        mode='text', auto_set=False, enter_set=True)

    pulse       = Range(low=1., high=100000.,   value=1000.,    desc='Length of the Pulse that is Calibrated',   label='pulse [ns]',    mode='text', auto_set=False, enter_set=True)
    delay       = Range(low=1., high=100000.,   value=100.,     desc='Delay between two pulses',   label='delay [ns]',    mode='text', auto_set=False, enter_set=True)

    N = Range(low=0, high=100,   value=6,    desc='Maximum Number of pulses',      label='N',       mode='text', auto_set=False, enter_set=True)

    def start_up(self):
        ha.PulseGenerator().Night()
        ha.Microwave().setOutput(self.power, self.frequency)

    def shut_down(self):
        ha.PulseGenerator().Light()
        ha.Microwave().setOutput(None, self.frequency)

    def generate_sequence(self):
        mw = self.switch
        laser = self.laser
        wait = self.wait
        delay = self.delay
        pulse = self.pulse
        N = self.N
        sequence = []
        for i in range(N):
            sequence += i * [  ([mw],pulse),  ([],delay)  ] + [  (['laser','aom'],laser),  ([],wait)  ]
        sequence += [  (['sequence'],100)  ]
        return sequence

    get_set_items = Pulsed.get_set_items + ['frequency','power','laser','wait','pulse','delay','N','switch']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),
                                                   Item('switch', style='custom', editor=EnumEditor(cols=3, values={'mw_x':'1:X', 'mw_y':'2:Y'},),)
                                                   ),
                                            HGroup(Item('pulse',         width=-80, enabled_when='state == "idle"'),
                                                   Item('delay',         width=-80, enabled_when='state == "idle"'),
                                                   Item('N',             width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'
                                            ),
                                     ),
                              ),
                              title='Pulse Calibration',
                              )

class Pi2PiCal( Pulsed ):
    
    """Pulse Calibration with initial and final pi/2_x pulses, pi_x and pi_y pulse sequences and bright / dark reference."""

    frequency   = Range(low=1,      high=20e9,  value=2.8705e9, desc='microwave frequency', label='frequency [Hz]', mode='text', auto_set=False, enter_set=True)
    power       = Range(low=-100.,  high=25.,   value=-20,      desc='microwave power',     label='power [dBm]',    mode='text', auto_set=False, enter_set=True)

    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)
    wait        = Range(low=1., high=100000.,   value=1000.,    desc='wait [ns]',       label='wait [ns]',        mode='text', auto_set=False, enter_set=True)

    t_pi2_x     = Range(low=0., high=100000.,   value=1000.,    desc='pi/2 pulse length (x)',   label='pi/2 x [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_y      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length (y)',     label='pi y [ns]',      mode='text', auto_set=False, enter_set=True)
    t_pi_x      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length (x)',     label='pi x [ns]',   mode='text', auto_set=False, enter_set=True)
    delay       = Range(low=1., high=100000.,   value=100.,     desc='Delay between two pulses',   label='delay [ns]',    mode='text', auto_set=False, enter_set=True)

    n_pi        = Range(low=0,  high=100,       value=10,       desc='number of pi pulses',   label='n pi',           mode='text', auto_set=False, enter_set=True)
    n_ref       = Range(low=1,  high=100,       value=10,       desc='number of reference pulses',   label='n ref',   mode='text', auto_set=False, enter_set=True)

    def start_up(self):
        ha.PulseGenerator().Night()
        ha.Microwave().setOutput(self.power, self.frequency)

    def shut_down(self):
        ha.PulseGenerator().Light()
        ha.Microwave().setOutput(None, self.frequency)

    def generate_sequence(self):

        # parameters        
        laser = self.laser
        wait  = self.wait
        t_pi2_x  = self.t_pi2_x
        t_pi_y   = self.t_pi_y
        t_pi_x   = self.t_pi_x
        n_pi     = self.n_pi
        n_ref    = self.n_ref
        t        = self.delay
        sequence = []
        
        # first sequence
        sequence +=     [ (['laser','aom'],laser),  ([],wait) ]
        for n in range(n_pi+1):
            sequence +=     [ (['mw_x'],t_pi2_x),       ([],t)    ]
            sequence += n * [ (['mw_x'],t_pi_x),        ([],2*t)  ] 
            sequence +=     [ (['laser','aom'],laser),  ([],wait) ]
        sequence +=     [ (['mw_x'],t_pi2_x),       ([],t)    ]
        sequence += n * [ (['mw_x'],t_pi_x),        ([],2*t)  ] 
        sequence +=     [ (['mw_x'],t_pi2_x),       ([],t)    ]
        sequence +=     [ (['laser','aom'],laser),  ([],wait) ]
        
        # second sequence
        sequence +=     [ (['laser','aom'],laser),  ([],wait) ]
        for n in range(n_pi+1):
            sequence +=     [ (['mw_x'],t_pi2_x),       ([],t)    ]
            sequence += n * [ (['mw_y'],t_pi_y),        ([],2*t)  ] 
            sequence +=     [ (['laser','aom'],laser),  ([],wait) ]
        sequence +=     [ (['mw_x'],t_pi2_x),       ([],t)    ]
        sequence += n * [ (['mw_y'],t_pi_y),        ([],2*t)  ]
        sequence +=     [ (['mw_x'],t_pi2_x),       ([],t)    ]
        sequence +=     [ (['laser','aom'],laser),  ([],wait) ]
        
        # bright state reference
        sequence += n_ref * [ (['laser','aom'],laser),  ([],wait) ]
        
        # dark state reference
        sequence += n_ref * [ (['mw_y'],t_pi_y),  (['laser','aom'],laser),  ([],wait)  ]

        # start trigger
        sequence += [ (['sequence'],100)  ]

        
        return sequence

    get_set_items = Pulsed.get_set_items + ['frequency','power','laser','wait','t_pi2_x','t_pi_x','t_pi_y','delay','n_pi','n_ref']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('t_pi2_x',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_y',        width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_x',        width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('delay',         width=-80, enabled_when='state == "idle"'),
                                                   Item('n_pi',          width=-80, enabled_when='state == "idle"'),
                                                   Item('n_ref',         width=-80, enabled_when='state == "idle"'),),
                                            label='manipulation'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='detection'
                                            ),
                                     ),
                              ),
                              title='Pi/2 Pi Calibration',
                              )

class Test( Pulsed ):
    
    """Test sequence."""

    frequency   = Range(low=1,      high=20e9,  value=2.8705e9, desc='microwave frequency', label='frequency [Hz]', mode='text', auto_set=False, enter_set=True)
    power       = Range(low=-100.,  high=25.,   value=-20,      desc='microwave power',     label='power [dBm]',    mode='text', auto_set=False, enter_set=True)

    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)
    wait        = Range(low=1., high=100000.,   value=1000.,    desc='wait [ns]',       label='wait [ns]',        mode='text', auto_set=False, enter_set=True)

    t_pi2_x     = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length (x)',   label='pi/2 x [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_x      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length (x)',     label='pi x [ns]',   mode='text', auto_set=False, enter_set=True)

    n           = Range(low=0,  high=100,       value=10,       desc='number of pulses',   label='n',           mode='text', auto_set=False, enter_set=True)

    def start_up(self):
        ha.PulseGenerator().Night()
        ha.Microwave().setOutput(self.power, self.frequency)

    def shut_down(self):
        ha.PulseGenerator().Light()
        ha.Microwave().setOutput(None, self.frequency)

    def generate_sequence(self):

        # parameters        
        laser = self.laser
        wait  = self.wait
        t_pi2_x  = self.t_pi2_x
        t_pi_x   = self.t_pi_x
        n        = self.n
        
        # start trigger
        sequence = [ (['sequence'],100)  ]

        # three sequence
        sequence += n * [ (['laser','aom'],laser),  ([],wait) ]
        sequence += n * [ (['mw_x'],t_pi_x),  (['laser','aom'],laser),  ([],wait)  ]
        sequence += n * [ (['mw_x'],t_pi2_x), (['laser','aom'],laser),  ([],wait)  ]
        
        return sequence

    get_set_items = Pulsed.get_set_items + ['frequency','power','laser','wait','t_pi2_x','t_pi_x','n']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('t_pi2_x',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_x',        width=-80, enabled_when='state == "idle"'),
                                                   Item('n',             width=-80, enabled_when='state == "idle"'),),
                                            label='manipulation'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='detection'
                                            ),
                                     ),
                              ),
                              title='Test bright, dark and pi/2 reproducability',
                              )

class TestBright( Pulsed ):
    
    """Test sequence."""

    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)
    wait        = Range(low=1., high=100000.,   value=1000.,    desc='wait [ns]',       label='wait [ns]',        mode='text', auto_set=False, enter_set=True)

    n           = Range(low=0,  high=100,       value=10,       desc='number of pulses',   label='n',           mode='text', auto_set=False, enter_set=True)

    def start_up(self):
        ha.PulseGenerator().Night()

    def shut_down(self):
        ha.PulseGenerator().Light()

    def generate_sequence(self):

        # parameters        
        laser = self.laser
        wait  = self.wait
        n        = self.n
        
        sequence = [ (['sequence'],100)  ] + n * [ (['laser','aom'],laser),  ([],wait) ]

        return sequence

    get_set_items = Pulsed.get_set_items + ['laser','wait','n']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('n',             width=-80, enabled_when='state == "idle"'),),
                                            label='manipulation'),
                                     VGroup(HGroup(Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='detection'
                                            ),
                                     ),
                              ),
                              title='Bright state detection',
                              )

class SingletDecay( Pulsed ):
    
    """Singlet Decay."""

    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)

    tau_begin   = Range(low=0., high=1e8,       value=0.,     desc='tau begin [ns]',  label='tau begin [ns]',   mode='text', auto_set=False, enter_set=True)
    tau_end     = Range(low=1., high=1e8,       value=300.,     desc='tau end [ns]',    label='tau end [ns]',     mode='text', auto_set=False, enter_set=True)
    tau_delta   = Range(low=1., high=1e6,       value=3.,      desc='delta tau [ns]',  label='delta tau [ns]',   mode='text', auto_set=False, enter_set=True)

    tau = Array( value=np.array((0.,1.)) )

    def apply_parameters(self):
        """Overwrites apply_parameters() from pulsed. Prior to generating sequence, etc., generate the tau mesh."""
        self.tau = np.arange(self.tau_begin, self.tau_end, self.tau_delta)
        Pulsed.apply_parameters(self)

    def start_up(self):
        ha.PulseGenerator().Night()

    def shut_down(self):
        ha.PulseGenerator().Light()

    def generate_sequence(self):

        tau = self.tau
        laser = self.laser

        sequence = [ (['sequence'],100) ]
        for t in tau:
            sequence += [  ([],t),  (['laser','aom'],laser)  ]
        return sequence

    get_set_items = Pulsed.get_set_items + ['tau_begin','tau_end','tau_delta','laser','tau']
    get_set_order = ['tau','time_bins','count_data']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='manipulation'),
                                     VGroup(HGroup(Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='detection'
                                            ),
                                     ),
                              ),
                              title='Singlet decay',
                              )
    
class SingletDecayDouble( Pulsed ):
    
    """Singlet Decay."""

    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)

    tau_begin   = Range(low=0., high=1e8,       value=0.,     desc='tau begin [ns]',  label='tau begin [ns]',   mode='text', auto_set=False, enter_set=True)
    tau_end     = Range(low=1., high=1e8,       value=300.,     desc='tau end [ns]',    label='tau end [ns]',     mode='text', auto_set=False, enter_set=True)
    tau_delta   = Range(low=1., high=1e6,       value=3.,      desc='delta tau [ns]',  label='delta tau [ns]',   mode='text', auto_set=False, enter_set=True)

    tau = Array( value=np.array((0.,1.)) )

    def apply_parameters(self):
        """Overwrites apply_parameters() from pulsed. Prior to generating sequence, etc., generate the tau mesh."""
        self.tau = np.arange(self.tau_begin, self.tau_end, self.tau_delta)
        Pulsed.apply_parameters(self)

    def start_up(self):
        ha.PulseGenerator().Night()

    def shut_down(self):
        ha.PulseGenerator().Light()

    def generate_sequence(self):

        tau = self.tau
        laser = self.laser

        sequence = [ (['sequence'],100) ]
        for t in tau:
            sequence += [  ([],t),  (['laser','aom'],laser)  ]
        for t in tau[::-1]:
            sequence += [  ([],t),  (['laser','aom'],laser)  ]
        return sequence

    get_set_items = Pulsed.get_set_items + ['tau_begin','tau_end','tau_delta','laser','tau']
    get_set_order = ['tau','time_bins','count_data']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='manipulation'),
                                     VGroup(HGroup(Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='detection'
                                            ),
                                     ),
                              ),
                              title='Singlet decay double',
                              )
    
class Hahn( Rabi ):
    
    """Defines a Hahn-Echo measurement."""

    t_pi2       = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length',   label='pi/2 [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi        = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length',     label='pi [ns]',      mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2  = self.t_pi2
        t_pi   = self.t_pi
        sequence = []
        for t in tau:
            sequence += [  (['mw'],t_pi2), ([],0.5*t), (['mw'],t_pi), ([],0.5*t), (['mw'],t_pi2), (['laser','aom'],laser), ([],wait)  ]
        sequence += [  (['sequence'],100)  ]
        return sequence

    get_set_items = Rabi.get_set_items + ['t_pi2','t_pi']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi2',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi',          width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'
                                            ),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'
                                            ),
                                     ),
                              ),
                       title='Hahn-Echo Measurement',
                       )
    

class Hahn3pi2( Rabi ):
    
    """Defines a Hahn-Echo measurement with both pi/2 and 3pi/2 readout pulse."""

    t_pi2       = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length',   label='pi/2 [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi        = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length',     label='pi [ns]',      mode='text', auto_set=False, enter_set=True)
    t_3pi2      = Range(low=1., high=100000.,   value=1000.,    desc='3pi/2 pulse length',  label='3pi/2 [ns]',   mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2  = self.t_pi2
        t_pi   = self.t_pi
        t_3pi2 = self.t_3pi2
        sequence = []
        for t in tau:
            sequence += [  (['mw'],t_pi2), ([],0.5*t), (['mw'],t_pi), ([],0.5*t), (['mw'],t_pi2), (['laser','aom'],laser), ([],wait)  ]
        for t in tau:
            sequence += [  (['mw'],t_pi2), ([],0.5*t), (['mw'],t_pi), ([],0.5*t), (['mw'],t_3pi2), (['laser','aom'],laser), ([],wait)  ]
        sequence += [  (['sequence'],100)  ]
        return sequence

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi2',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi',          width=-80, enabled_when='state == "idle"'),
                                                   Item('t_3pi2',        width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'
                                            ),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'
                                            ),
                                     ),
                              ),
                       title='Hahn-Echo Measurement with both pi/2 and 3pi/2 readout pulse',
                       )
    
    get_set_items = Rabi.get_set_items + ['t_pi2','t_pi','t_3pi2']

class CPMG3pi2( Rabi ):
    
    """
    Defines a CPMG measurement with both pi/2 and 3pi/2 readout pulse,
    using a second microwave switch for 90 degree phase shifted pi pulses.
    
    Includes also bright (no pulse) and dark (pi_y pulse) reference points.
    """

    t_pi2_x     = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length (x)',   label='pi/2 x [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_y      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length (y)',     label='pi y [ns]',      mode='text', auto_set=False, enter_set=True)
    t_3pi2_x    = Range(low=1., high=100000.,   value=1000.,    desc='3pi/2 pulse length (x)',  label='3pi/2 x [ns]',   mode='text', auto_set=False, enter_set=True)
    n_pi        = Range(low=1,  high=100,       value=5,        desc='number of pi pulses',   label='n pi',           mode='text', auto_set=False, enter_set=True)
    n_ref       = Range(low=1,  high=100,       value=10,       desc='number of reference pulses',   label='n ref',   mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2_x  = self.t_pi2_x
        t_pi_y   = self.t_pi_y
        t_3pi2_x = self.t_3pi2_x
        n_pi     = self.n_pi
        n_ref    = self.n_ref
        sequence = []
        for t in tau:
            sequence +=            [ (['mw_x'],t_pi2_x),       ([],t/float(2*n_pi))  ]
            sequence += (n_pi-1) * [ (['mw_y'],t_pi_y),        ([],t/float(n_pi))    ] 
            sequence +=            [ (['mw_y'],t_pi_y),        ([],t/float(2*n_pi))  ]
            sequence +=            [ (['mw_x'],t_pi2_x)                              ]
            sequence +=            [ (['laser','aom'],laser),  ([],wait)             ]
        for t in tau:
            sequence +=             [ (['mw_x'],t_pi2_x),       ([],t/float(2*n_pi))  ]
            sequence += (n_pi-1) *  [ (['mw_y'],t_pi_y),        ([],t/float(n_pi))    ] 
            sequence +=             [ (['mw_y'],t_pi_y),        ([],t/float(2*n_pi))  ]
            sequence +=             [ (['mw_x'],t_3pi2_x)                             ]
            sequence +=             [ (['laser','aom'],laser),  ([],wait)             ]            
        sequence +=    n_ref * [ (['laser','aom'],laser),  ([],wait)       ]
        sequence +=    n_ref * [ (['mw_y'],t_pi_y),  (['laser','aom'],laser),  ([],wait)      ]
        sequence +=            [ (['sequence'],100)                        ]
        return sequence

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('t_pi2_x',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_y',        width=-80, enabled_when='state == "idle"'),
                                                   Item('t_3pi2_x',      width=-80, enabled_when='state == "idle"'),
                                                   Item('n_pi',          width=-80, enabled_when='state == "idle"'),
                                                   Item('n_ref',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'
                                            ),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'
                                            ),
                                     ),
                              ),
                       title='CPMG Measurement with both pi/2 and 3pi/2 readout pulse',
                       )
    
    get_set_items = Rabi.get_set_items + ['t_pi2_x','t_pi_y','t_3pi2_x','n_pi']

class CPMG( Rabi ):
    
    """Defines a basic CPMG measurement with a single sequence and bright / dark reference points."""

    t_pi2_x     = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 x pulse length',   label='pi/2 x [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_y      = Range(low=1., high=100000.,   value=1000.,    desc='pi y pulse length',     label='pi y [ns]',      mode='text', auto_set=False, enter_set=True)
    n_pi        = Range(low=1,  high=100,       value=5,        desc='number of pi pulses',   label='n pi',           mode='text', auto_set=False, enter_set=True)
    n_ref       = Range(low=1,  high=100,       value=10,       desc='number of reference pulses',   label='n ref',   mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2_x = self.t_pi2_x
        t_pi_y  = self.t_pi_y
        n_pi    = self.n_pi
        n_ref   = self.n_ref
        sequence = []
        for t in tau:
            sequence +=            [ (['mw_x'],t_pi2_x),       ([],t/float(2*n_pi)) ]
            sequence += (n_pi-1) * [ (['mw_y'],t_pi_y),        ([],t/float(n_pi))   ] 
            sequence +=            [ (['mw_y'],t_pi_y),        ([],t/float(2*n_pi)) ]
            sequence +=            [ (['mw_x'],t_pi2_x)                             ]
            sequence +=            [ (['laser','aom'],laser),  ([],wait)            ]
        sequence +=    n_ref * [ (['laser','aom'],laser),  ([],wait)       ]
        sequence +=    n_ref * [ (['mw_y'],t_pi_y),  (['laser','aom'],laser),  ([],wait)      ]
        sequence +=            [ (['sequence'],100)                        ]
        return sequence

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),),       
                                            HGroup(Item('t_pi2_x',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_y',        width=-80, enabled_when='state == "idle"'),
                                                   Item('n_pi',          width=-80, enabled_when='state == "idle"'),
                                                   Item('n_ref',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'),
                                     ),
                              ),
                       title='Basic CPMG Measurement with pi/2 pulse on B and pi pulses on A',
                       )
    
    get_set_items = Rabi.get_set_items + ['t_pi2_x','t_pi_y','n_pi']


class CPMGxy( Rabi ):
    
    """
    Defines a CPMG measurement with both pi_x and pi_y pulses.

    Includes also bright (no pulse) and dark (pi_y pulse) reference points.
    """

    t_pi2_x     = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length (x)',   label='pi/2 x [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_x      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length (x)',     label='pi x [ns]',      mode='text', auto_set=False, enter_set=True)
    t_pi_y      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length (y, 90 degree)',  label='pi y [ns]',   mode='text', auto_set=False, enter_set=True)
    n_pi        = Range(low=1,  high=100,       value=5,        desc='number of pi pulses',   label='n pi',           mode='text', auto_set=False, enter_set=True)
    n_ref       = Range(low=1,  high=100,       value=10,       desc='number of reference pulses',   label='n ref',   mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2_x  = self.t_pi2_x
        t_pi_x   = self.t_pi_x
        t_pi_y = self.t_pi_y
        n_pi   = self.n_pi
        n_ref    = self.n_ref
        sequence = []
        for t in tau:
            sequence +=             [ (['mw_x'],t_pi2_x),       ([],t/float(2*n_pi))  ]
            sequence += (n_pi-1) *  [ (['mw_x'],t_pi_x),        ([],t/float(n_pi))    ] 
            sequence +=             [ (['mw_x'],t_pi_x),        ([],t/float(2*n_pi))  ]
            sequence +=             [ (['mw_x'],t_pi2_x)                              ]
            sequence +=             [ (['laser','aom'],laser),  ([],wait)             ]
        for t in tau:
            sequence +=             [ (['mw_x'],t_pi2_x),       ([],t/float(2*n_pi))  ]
            sequence += (n_pi-1) *  [ (['mw_y'],t_pi_y),        ([],t/float(n_pi))    ] 
            sequence +=             [ (['mw_y'],t_pi_y),        ([],t/float(2*n_pi))  ]
            sequence +=             [ (['mw_x'],t_pi2_x)                              ]
            sequence +=             [ (['laser','aom'],laser),  ([],wait)             ]            
        sequence +=    n_ref * [ (['laser','aom'],laser),  ([],wait)        ]
        sequence +=    n_ref * [ (['mw_y'],t_pi_y),  (['laser','aom'],laser),  ([],wait)      ]
        sequence +=            [ (['sequence'],100)                         ]
        return sequence

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('t_pi2_x',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_x',        width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_y',        width=-80, enabled_when='state == "idle"'),
                                                   Item('n_pi',          width=-80, enabled_when='state == "idle"'),
                                                   Item('n_ref',         width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'),
                                     ),
                              ),
                       title='CPMG Measurement with both pix and piys pulses',
                       )
    
    get_set_items = Rabi.get_set_items + ['t_pi2_x','t_pi_x','t_pi_y','n_pi']

class T1pi( Rabi ):
    
    """Defines a T1 measurement with pi pulse."""

    t_pi        = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length',     label='pi [ns]',      mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi   = self.t_pi
        sequence = []
        for t in tau:
            sequence.append(  ([       ],   t )  )
            sequence.append(  (['laser','aom'],  laser  )  )
            sequence.append(  ([       ],  wait   )  )
        for t in tau:
            sequence.append(  (['mw'   ],  t_pi   )  )
            sequence.append(  ([       ],   t )  )
            sequence.append(  (['laser','aom'],  laser  )  )
            sequence.append(  ([       ],  wait   )  )
        sequence.append(  (['sequence'], 100  )  )
        return sequence

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),                                                
                                                   Item('t_pi',          width=-80, enabled_when='state == "idle"'),),                                                      
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'),
                                     ),
                              ),
                       title='T1pi',
                       )
    
    get_set_items = Rabi.get_set_items + ['t_pi']

class FID3pi2( Rabi ):
    
    """Defines a FID measurement with both pi/2 and 3pi/2 readout pulse."""

    t_pi2       = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length',   label='pi/2 [ns]',    mode='text', auto_set=False, enter_set=True)
    t_3pi2      = Range(low=1., high=100000.,   value=1000.,    desc='3pi/2 pulse length',  label='3pi/2 [ns]',   mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2  = self.t_pi2
        t_3pi2 = self.t_3pi2
        sequence = []
        for t in tau:
            sequence.append(  (['mw'   ],            t_pi2  )  )
            sequence.append(  ([       ],            t      )  )
            sequence.append(  (['mw'   ],            t_pi2  )  )
            sequence.append(  (['laser','aom'],            laser  )  )
            sequence.append(  ([       ],            wait   )  )
        for t in tau:
            sequence.append(  (['mw'   ],            t_pi2  )  )
            sequence.append(  ([       ],            t      )  )
            sequence.append(  (['mw'   ],            t_3pi2 )  )
            sequence.append(  (['laser','aom'],            laser  )  )
            sequence.append(  ([       ],            wait   )  )
        sequence.append(  (['sequence'], 100  )  )
        return sequence
    
    #items to be saved
    get_set_items = Rabi.get_set_items + ['t_pi2','t_3pi2']

    # gui elements
    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi2',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_3pi2',        width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'),
                                     ),
                              ),
                       title='FID3pi2',
                       )
    


class HardDQTFID( Rabi ):
    
    """Defines a Double Quantum Transition FID Measurement with pi-tau-pi."""

    t_pi        = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length',     label='pi [ns]',      mode='text', auto_set=False, enter_set=True)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi   = self.t_pi
        sequence = []
        for t in tau:
            sequence.append(  (['mw'   ],            t_pi   )  )
            sequence.append(  ([       ],            t      )  )
            sequence.append(  (['mw'   ],            t_pi   )  )
            sequence.append(  (['laser','aom'],            laser  )  )
            sequence.append(  ([       ],            wait   )  )
        sequence.append(  (['sequence'], 100  )  )
        return sequence

    get_set_items = Rabi.get_set_items + ['t_pi']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                   Item('power',         width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi',          width=-80, enabled_when='state == "idle"'),),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                            label='settings'),
                                     ),
                              ),
                       title='HardDQTFID',
                       )

"""
class HardDQTFIDTauMod( HardDQTFID ):
    tau_gen = Code("tg = np.arange(self.tau_begin, self.tau_end, self.tau_delta)\ntf = np.arange(0., 500., 50.)\ntau = np.array(())\nfor t0 in tg:\n   tau=np.append(tau,tf+t0)\nself.tau=tau")
    def start_up(self):
        # this needs to come first, because Pulsed.start_up() will
        # generate the sequence and thereby call the generate_sequence method of RabiMeasurement
        # and this method needs the correct tau 
        exec self.tau_gen
        Pulsed.start_up(self)
        ha.PulseGenerator().Night()
        ha.Microwave().setOutput(self.power, self.frequency)
        
    traits_view = View( VGroup( HGroup(Item('state', show_label=False, style='custom',
                                            editor=EnumEditor(values={'idle':'1:idle','run':'2:run',},cols=2),),),
                                Tabbed( VGroup(HGroup(Item('frequency',     width=-80, enabled_when='state == "idle"'),
                                                      Item('power',         width=-80, enabled_when='state == "idle"'),
                                                      Item('t_pi',          width=-80, enabled_when='state == "idle"'),),
                                               HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                      Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                      Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                               label='parameter'),
                                        VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                      Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                      Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                      Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                               label='settings'),
                                        VGroup(HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                      Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                      Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                                Item('tau_gen',   width=-200, height=-200, enabled_when='state == "idle"'),
                                               label='taumod'),
                                            
                                ),
                        ),
                        title='HardDQTFIDTauMod',
                  )
"""

class DQTFID3pi2( Pulsed ):
    
    """Defines a pulsed measurement with two microwave transitions and FID sequence including 3pi2 readout pulses.TODO: untested"""
    frequency_a = Range(low=1,      high=20e9,  value=2.8705e9, desc='microwave frequency A',   label='frequency A [Hz]', mode='text', auto_set=False, enter_set=True)
    frequency_b = Range(low=1,      high=20e9,  value=2.8705e9, desc='microwave frequency B',   label='frequency B [Hz]', mode='text', auto_set=False, enter_set=True)
    power_a     = Range(low=-100.,  high=25.,   value=-20,      desc='microwave power A',       label='power A [dBm]',    mode='text', auto_set=False, enter_set=True)
    power_b     = Range(low=-100.,  high=25.,   value=-20,      desc='microwave power B',       label='power B [dBm]',    mode='text', auto_set=False, enter_set=True)

    tau_begin   = Range(low=1., high=1e8,       value=10.0,     desc='tau begin [ns]',  label='tau begin [ns]',   mode='text', auto_set=False, enter_set=True)
    tau_end     = Range(low=1., high=1e8,       value=500.,     desc='tau end [ns]',    label='tau end [ns]',     mode='text', auto_set=False, enter_set=True)
    tau_delta   = Range(low=1., high=1e6,       value=50.,      desc='delta tau [ns]',  label='delta tau [ns]',   mode='text', auto_set=False, enter_set=True)
    laser       = Range(low=1., high=100000.,   value=3000.,    desc='laser [ns]',      label='laser [ns]',       mode='text', auto_set=False, enter_set=True)
    wait        = Range(low=1., high=100000.,   value=1000.,    desc='wait [ns]',       label='wait [ns]',        mode='text', auto_set=False, enter_set=True)

    t_pi2_a     = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length A',     label='pi/2 A [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_a      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length A',       label='pi A [ns]',      mode='text', auto_set=False, enter_set=True)
    t_3pi2_a    = Range(low=1., high=100000.,   value=1000.,    desc='3pi/2 pulse length A',    label='3pi/2 A [ns]',   mode='text', auto_set=False, enter_set=True)
    t_pi2_b     = Range(low=1., high=100000.,   value=1000.,    desc='pi/2 pulse length B',     label='pi/2 B [ns]',    mode='text', auto_set=False, enter_set=True)
    t_pi_b      = Range(low=1., high=100000.,   value=1000.,    desc='pi pulse length B',       label='pi B [ns]',      mode='text', auto_set=False, enter_set=True)
    t_3pi2_b    = Range(low=1., high=100000.,   value=1000.,    desc='3pi/2 pulse length B',    label='3pi/2 B [ns]',   mode='text', auto_set=False, enter_set=True)

    tau = Array( value=np.array((0.,1.)) )

    def start_up(self):
        self.tau = np.arange(self.tau_begin, self.tau_end, self.tau_delta)
        ha.PulseGenerator().Night()
        ha.MicrowaveA().setOutput(self.power_a, self.frequency_a)
        ha.MicrowaveB().setOutput(self.power_b, self.frequency_b)

    def shut_down(self):
        ha.PulseGenerator().Light()
        ha.MicrowaveA().setOutput(None, self.frequency_a)
        ha.MicrowaveB().setOutput(None, self.frequency_b)

    def generate_sequence(self):
        tau = self.tau
        laser = self.laser
        wait  = self.wait
        t_pi2_a  = self.t_pi2_a
        t_pi_a   = self.t_pi_a
        t_3pi2_a = self.t_3pi2_a
        t_pi2_b  = self.t_pi2_b
        t_pi_b   = self.t_pi_b
        t_3pi2_b = self.t_3pi2_b
        sequence = []
        for t in tau:
            sequence.append(  (['mw_a'   ],  t_pi2_a  )  )
            sequence.append(  (['mw_b'   ],  t_pi_b   )  )
            sequence.append(  ([         ],  t        )  )
            sequence.append(  (['mw_b'   ],  t_pi_b   )  )
            sequence.append(  (['mw_a'   ],  t_pi2_a  )  )
            sequence.append(  (['laser','aom'],  laser    )  )
            sequence.append(  ([         ],  wait     )  )
        for t in tau:
            sequence.append(  (['mw_a'   ],  t_pi2_a  )  )
            sequence.append(  (['mw_b'   ],  t_pi_b   )  )
            sequence.append(  ([         ],  t        )  )
            sequence.append(  (['mw_b'   ],  t_pi_b   )  )
            sequence.append(  (['mw_a'   ], t_3pi2_a  )  )
            sequence.append(  (['laser','aom'],  laser    )  )
            sequence.append(  ([         ],  wait     )  )
        sequence.append(  (['sequence'], 100  )  )
        return sequence

    get_set_items = Pulsed.get_set_items + ['frequency_a','frequency_b','power_a','power_b',
                                                       't_pi2_a', 't_pi_a', 't_3pi2_a', 't_pi2_b', 't_pi_b', 't_3pi2_b',
                                                       'tau_begin','tau_end','tau_delta','laser','wait','tau']

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Tabbed(VGroup(HGroup(Item('frequency_a',   width=-80, enabled_when='state == "idle"'),
                                                   Item('power_a',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi2_a',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_a',        width=-80, enabled_when='state == "idle"'),
                                                   Item('t_3pi2_a',      width=-80, enabled_when='state == "idle"'),
                                            ),
                                            HGroup(Item('frequency_b',   width=-80, enabled_when='state == "idle"'),
                                                   Item('power_b',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi2_b',       width=-80, enabled_when='state == "idle"'),
                                                   Item('t_pi_b',        width=-80, enabled_when='state == "idle"'),
                                                   Item('t_3pi2_b',      width=-80, enabled_when='state == "idle"'),
                                            ),
                                            HGroup(Item('tau_begin',     width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_end',       width=-80, enabled_when='state == "idle"'),
                                                   Item('tau_delta',     width=-80, enabled_when='state == "idle"'),),
                                            label='parameter'),
                                     VGroup(HGroup(Item('laser',         width=-80, enabled_when='state == "idle"'),
                                                   Item('wait',          width=-80, enabled_when='state == "idle"'),
                                                   Item('record_length', width=-80, enabled_when='state == "idle"'),
                                                   Item('bin_width',     width=-80, enabled_when='state == "idle"'),),
                                                   label='settings'),
                                     ),
                              ),
                       title='DQTFID3pi2 Measurement',
                       )


if __name__ == '__main__':
    
    logging.getLogger().addHandler(logging.StreamHandler())
    logging.getLogger().setLevel(logging.DEBUG)
    logging.getLogger().info('Starting logger.')
    
    from emod import JobManager
    
    JobManager().start()
    
    r1 = Rabi()
    r2 = Rabi()
    
    r1.edit_traits()
    r2.edit_traits()
    